package websocket

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"sync"

	"github.com/gorilla/websocket"
	"monitor/models"
)

var (
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool {
			return true // 允许所有来源的WebSocket连接
		},
	}

	// 全局事件广播管理器
	broadcastManager = NewBroadcastManager()
)

// Client 表示一个WebSocket客户端连接
type Client struct {
	conn *websocket.Conn
	mutex sync.Mutex
}

// BroadcastManager 管理所有WebSocket客户端连接
type BroadcastManager struct {
	clients    map[*Client]bool
	mutex      sync.RWMutex
	register   chan *Client
	unregister chan *Client
}

// NewBroadcastManager 创建一个新的广播管理器
func NewBroadcastManager() *BroadcastManager {
	bm := &BroadcastManager{
		clients:    make(map[*Client]bool),
		register:   make(chan *Client),
		unregister: make(chan *Client),
	}
	go bm.run()
	return bm
}

// run 运行广播管理器的主循环
func (bm *BroadcastManager) run() {
	for {
		select {
		case client := <-bm.register:
			bm.mutex.Lock()
			bm.clients[client] = true
			bm.mutex.Unlock()
		case client := <-bm.unregister:
			bm.mutex.Lock()
			if _, ok := bm.clients[client]; ok {
				delete(bm.clients, client)
				client.conn.Close()
			}
			bm.mutex.Unlock()
		}
	}
}

// BroadcastEvent 向所有连接的客户端广播事件
func (bm *BroadcastManager) BroadcastEvent(event *models.Event) {
	bm.mutex.RLock()
	defer bm.mutex.RUnlock()

	eventJSON, err := json.Marshal(event)
	if err != nil {
		log.Printf("事件序列化失败: %v\n", err)
		return
	}

	for client := range bm.clients {
		client.mutex.Lock()
		err := client.conn.WriteMessage(websocket.TextMessage, eventJSON)
		client.mutex.Unlock()

		if err != nil {
			log.Printf("发送消息失败: %v\n", err)
			bm.unregister <- client
		}
	}
}

// GetBroadcastManager 获取全局广播管理器实例
func GetBroadcastManager() *BroadcastManager {
	return broadcastManager
}

// HandleWebSocket 处理WebSocket连接请求
func HandleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket升级失败: %v\n", err)
		return
	}

	client := &Client{conn: conn}
	broadcastManager.register <- client

	// 处理连接关闭
	go func() {
		for {
			_, _, err := conn.ReadMessage()
			if err != nil {
				broadcastManager.unregister <- client
				break
			}
		}
	}()
}

// StartWebSocketServer 启动WebSocket服务器
func StartWebSocketServer(port int) error {
	http.HandleFunc("/ws", HandleWebSocket)
	addr := fmt.Sprintf(":%d", port)
	log.Printf("WebSocket服务器启动在 %s\n", addr)
	return http.ListenAndServe(addr, nil)
}